#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <string.h>
#include <errno.h>
#include <signal.h>
#include <poll.h>
#include <time.h>

/* 包含驱动头文件中的定义 */
#include "../rk-pcie-ep.h"

#define DEVICE_PREFIX "/dev/pcie-rkep"
#define TEST_PATTERN 0xA5A5A5A5
#define TEST_SIZE 4096
#define MMAP_SIZE (4 * 1024 * 1024) /* 4MB */

static int fd;
static volatile int irq_received = 0;

/* 信号处理函数 */
void sigio_handler(int sig)
{
    irq_received = 1;
    printf("Received SIGIO signal!\n");
}

/* 测试MMAP内存区域 */
int test_mmap(void)
{
    void *mapped_addr;
    unsigned int *test_ptr;
    int i, ret = 0;

    printf("\n===== Testing MMAP =====\n");
    
    /* 映射内存区域 */
    mapped_addr = mmap(NULL, MMAP_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
    if (mapped_addr == MAP_FAILED) {
        perror("mmap failed");
        return -1;
    }
    printf("Mapped address: %p\n", mapped_addr);

    /* 写入测试模式 */
    test_ptr = (unsigned int *)mapped_addr;
    for (i = 0; i < TEST_SIZE / sizeof(int); i++) {
        test_ptr[i] = TEST_PATTERN + i;
    }
    printf("Written test pattern to MMAP area\n");

    /* 验证读取 */
    for (i = 0; i < TEST_SIZE / sizeof(int); i++) {
        if (test_ptr[i] != (TEST_PATTERN + i)) {
            printf("Verification failed at offset %d: expected 0x%X, got 0x%X\n",
                   i * 4, TEST_PATTERN + i, test_ptr[i]);
            ret = -1;
            break;
        }
    }
    
    if (ret == 0) {
        printf("MMAP read/write verification successful!\n");
    }

    /* 清理 */
    munmap(mapped_addr, MMAP_SIZE);
    return ret;
}

/* 测试BAR0读写 */
int test_bar0_rw(void)
{
    unsigned int *buffer;
    ssize_t bytes;
    int i, ret = 0;

    printf("\n===== Testing BAR0 Read/Write =====\n");
    
    /* 分配测试缓冲区 */
    buffer = malloc(TEST_SIZE);
    if (!buffer) {
        perror("malloc failed");
        return -1;
    }

    /* 准备测试数据 */
    for (i = 0; i < TEST_SIZE / sizeof(int); i++) {
        buffer[i] = 0xDEADBEEF + i;
    }

    /* 写入BAR0 */
    bytes = write(fd, buffer, TEST_SIZE);
    if (bytes != TEST_SIZE) {
        perror("write BAR0 failed");
        free(buffer);
        return -1;
    }
    printf("Wrote %zd bytes to BAR0\n", bytes);

    /* 清除缓冲区 */
    memset(buffer, 0, TEST_SIZE);

    /* 从BAR0读取 */
    lseek(fd, 0, SEEK_SET);
    bytes = read(fd, buffer, TEST_SIZE);
    if (bytes != TEST_SIZE) {
        perror("read BAR0 failed");
        free(buffer);
        return -1;
    }
    printf("Read %zd bytes from BAR0\n", bytes);

    /* 验证数据 */
    for (i = 0; i < TEST_SIZE / sizeof(int); i++) {
        if (buffer[i] != (0xDEADBEEF + i)) {
            printf("Verification failed at offset %d: expected 0x%X, got 0x%X\n",
                   i * 4, 0xDEADBEEF + i, buffer[i]);
            ret = -1;
            break;
        }
    }
    
    if (ret == 0) {
        printf("BAR0 read/write verification successful!\n");
    }

    free(buffer);
    return ret;
}

/* 测试IOCTL命令 */
int test_ioctl_commands(void)
{
    struct pcie_ep_dma_cache_cfg cache_cfg;
    unsigned long long phys_addr;
    int ret;

    printf("\n===== Testing IOCTL Commands =====\n");
    
    /* 获取物理地址 */
    ret = ioctl(fd, 0x4, &phys_addr); /* 自定义命令0x4 */
    if (ret) {
        perror("ioctl get physical address failed");
        return ret;
    }
    printf("Got physical address: 0x%llx\n", (unsigned long long)phys_addr);

    /* 测试缓存无效命令 */
    cache_cfg.addr = phys_addr;
    cache_cfg.size = TEST_SIZE;
    ret = ioctl(fd, PCIE_DMA_CACHE_INVALIDE, &cache_cfg);
    if (ret) {
        perror("ioctl cache invalidate failed");
        return ret;
    }
    printf("Cache invalidate successful\n");

    /* 测试缓存刷新命令 */
    ret = ioctl(fd, PCIE_DMA_CACHE_FLUSH, &cache_cfg);
    if (ret) {
        perror("ioctl cache flush failed");
        return ret;
    }
    printf("Cache flush successful\n");

    return 0;
}

/* 测试中断通知 */
int test_irq_notification(void)
{
    struct timespec ts;
    int oflags, ret;

    printf("\n===== Testing IRQ Notification =====\n");
    
    /* 设置信号处理 */
    signal(SIGIO, sigio_handler);

    /* 设置异步通知 */
    fcntl(fd, F_SETOWN, getpid());
    oflags = fcntl(fd, F_GETFL);
    fcntl(fd, F_SETFL, oflags | FASYNC);

    /* 触发中断 - 通过写入BAR0特定位置 */
    unsigned int cmd = 0x524b4501; /* RKEP_CMD_LOADER_RUN */
    lseek(fd, 0x400, SEEK_SET);
    if (write(fd, &cmd, sizeof(cmd)) != sizeof(cmd)) {
        perror("trigger command write failed");
        return -1;
    }
    printf("Triggered interrupt command\n");

    /* 等待中断通知 */
    printf("Waiting for interrupt (timeout 5s)...\n");
    irq_received = 0;
    ts.tv_sec = time(NULL) + 5;
    ts.tv_nsec = 0;
    
    while (!irq_received) {
        struct timespec now;
        clock_gettime(CLOCK_REALTIME, &now);
        if (now.tv_sec > ts.tv_sec || 
            (now.tv_sec == ts.tv_sec && now.tv_nsec > ts.tv_nsec)) {
            printf("Timeout waiting for interrupt!\n");
            return -1;
        }
        usleep(100000); /* 100ms */
    }

    printf("IRQ notification test successful!\n");
    return 0;
}

int main(int argc, char *argv[])
{
    char dev_path[256];
    int ret = 0;
    
    if (argc > 1) {
        snprintf(dev_path, sizeof(dev_path), "%s", argv[1]);
    } else {
        /* 自动检测设备文件 */
        snprintf(dev_path, sizeof(dev_path), "%s-0000:09:00.0", DEVICE_PREFIX);
    }

    printf("Opening PCIe EP device: %s\n", dev_path);
    fd = open(dev_path, O_RDWR);
    if (fd < 0) {
        perror("open device failed");
        return EXIT_FAILURE;
    }

    test_mmap();
    test_bar0_rw();
    test_ioctl_commands();
    test_irq_notification();

    close(fd);
    return ret ? EXIT_FAILURE : EXIT_SUCCESS;
}
